package me.zhengjie.config.sftp;

import com.jcraft.jsch.ChannelSftp;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.pool2.impl.GenericObjectPool;

import java.io.IOException;
import java.io.InputStream;

@Slf4j
public class SftpInputStreamWrapper extends InputStream {
    private final InputStream delegate;
    private final ChannelSftp sftp;
    private final GenericObjectPool<ChannelSftp> pool;
    private volatile boolean closed = false;

    public SftpInputStreamWrapper(InputStream delegate,
                                  ChannelSftp sftp,
                                  GenericObjectPool<ChannelSftp> pool) {
        this.delegate = delegate;
        this.sftp = sftp;
        this.pool = pool;
    }

    // 代理所有InputStream方法
    @Override
    public int read() throws IOException {
        return delegate.read();
    }

    @Override
    public int read(byte[] b) throws IOException {
        return delegate.read(b);
    }

    @Override
    public int read(byte[] b, int off, int len) throws IOException {
        return delegate.read(b, off, len);
    }

    @Override
    public long skip(long n) throws IOException {
        return delegate.skip(n);
    }

    @Override
    public int available() throws IOException {
        return delegate.available();
    }

    @Override
    public synchronized void mark(int readlimit) {
        delegate.mark(readlimit);
    }

    @Override
    public synchronized void reset() throws IOException {
        delegate.reset();
    }

    @Override
    public boolean markSupported() {
        return delegate.markSupported();
    }

    // 关键重写：关闭时归还连接
    @Override
    public void close() throws IOException {
        if (closed) return;

        try {
            delegate.close(); // 关闭文件流
        } finally {
            closed = true;
            try {
                if (sftp.isConnected()) {
                    pool.returnObject(sftp); // 归还连接至池
                } else {
                    pool.invalidateObject(sftp); // 连接已断开则标记失效
                }
            } catch (Exception e) {
                log.error("归还SFTP连接失败", e);
            }
        }
    }
}
